# 运行步骤：
# 1. pip install ffmpeg-python pydub noisereduce librosa numpy
# 2. python enhance_video.py

import ffmpeg
import numpy as np
from pydub import AudioSegment, silence
import noisereduce as nr
import os

def enhance_video_quality(input_path, output_path):
    """视频增强逻辑（保持不变）"""
    try:
        (
            ffmpeg
            .input(input_path)
            .output(output_path, vf='hqdn3d=4:3:6:4, unsharp=5:5:1.0:5:5:0.0', crf=18, preset='slow')
            .overwrite_output()
            .run()
        )
        return True
    except ffmpeg.Error as e:
        print("FFmpeg视频处理出错:", e.stderr)
        return False

def process_audio(input_audio_path, output_audio_path):
    """改进的音频处理流程"""
    # 加载音频并转为单声道
    audio = AudioSegment.from_wav(input_audio_path).set_channels(1)
    samples = np.array(audio.get_array_of_samples())
    rate = audio.frame_rate

    # 检测静音段（更严格的阈值）
    silence_ranges = silence.detect_silence(
        audio,
        min_silence_len=1000,
        silence_thresh=-45  # 更敏感的阈值
    )

    if not silence_ranges:
        print("警告：未检测到静音段，将尝试全段降噪")
        silence_ranges = [(0, len(audio))]

    # 自动选择最佳噪声样本（取最长静音段）
    noise_sample = max(silence_ranges, key=lambda x: x[1]-x[0])
    start, end = noise_sample
    noise_audio = audio[start:end]
    noise_samples = np.array(noise_audio.get_array_of_samples()).astype(np.float32) / (2**15)

    # 全局基础降噪
    audio_float = samples.astype(np.float32) / (2**15)
    global_denoised = nr.reduce_noise(
        y=audio_float,
        y_noise=noise_samples,
        sr=rate,
        prop_decrease=0.6  # 中等降噪强度
    )

    # 静音段增强降噪
    for seg in silence_ranges:
        seg_start = int(seg[0] * rate / 1000)
        seg_end = int(seg[1] * rate / 1000)
        segment = global_denoised[seg_start:seg_end]
        
        # 更激进的降噪参数
        enhanced_segment = nr.reduce_noise(
            y=segment,
            y_noise=noise_samples,
            sr=rate,
            prop_decrease=1.0,  # 最大降噪强度
            n_fft=1024,        # 更精细的频谱分析
            use_tqdm=False
        )
        global_denoised[seg_start:seg_end] = enhanced_segment

    # 转换为音频并保存
    processed_audio = AudioSegment(
        (global_denoised * (2**15)).astype(np.int16).tobytes(),
        frame_rate=rate,
        sample_width=2,
        channels=1
    )
    processed_audio.export(output_audio_path, format='wav')
    return True

def main(input_video_path, output_video_path):
    """完整的视频处理主流程"""
    try:
        # 创建临时文件路径
        temp_dir = "temp_processing"
        os.makedirs(temp_dir, exist_ok=True)
        
        temp_audio = os.path.join(temp_dir, "temp_audio.wav")
        temp_video = os.path.join(temp_dir, "temp_video.mp4")
        enhanced_audio = os.path.join(temp_dir, "enhanced_audio.wav")
        
        # 1. 提取音频
        print("正在提取音频...")
        (
            ffmpeg
            .input(input_video_path)
            .output(temp_audio, acodec='pcm_s16le', ac=1)
            .overwrite_output()
            .run(capture_stdout=True, capture_stderr=True)
        )
        
        # 2. 处理音频
        print("正在增强音频质量...")
        process_audio(temp_audio, enhanced_audio)
        
        # 3. 处理视频
        print("正在增强视频质量...")
        enhance_video_quality(input_video_path, temp_video)
        
        # 4. 合并处理后的视频和音频
        print("正在合并音视频...")
        stream = ffmpeg.input(temp_video)
        audio = ffmpeg.input(enhanced_audio)
        (
            ffmpeg
            .output(stream, audio, output_video_path, acodec='aac', vcodec='copy')
            .overwrite_output()
            .run(capture_stdout=True, capture_stderr=True)
        )
        
        # 5. 清理临时文件
        print("正在清理临时文件...")
        for temp_file in [temp_audio, temp_video, enhanced_audio]:
            if os.path.exists(temp_file):
                os.remove(temp_file)
        os.rmdir(temp_dir)
        
        print(f"处理完成！输出文件：{output_video_path}")
        return True
        
    except Exception as e:
        print(f"处理过程中出现错误：{str(e)}")
        return False

if __name__ == "__main__":
    input_video = "input.mp4"
    output_video = "output_enhanced.mp4"
    main(input_video, output_video)